﻿#include "Riemann2D.h"
#include "Timer.h"
#include "Application.h"
#include "EulerEquation.h"
#include "GeneratedMesh.h"
#include "Timer.h"
#include "Transient.h"

struct Riemann2DIC : WorkletInitCondition<Riemann2DIC>
{
  VTKM_EXEC ValueType ComputeValue(const Vec3f& p) const
  {
    ValueType primary_value;
    auto x = p[0];
    auto y = p[1];
    if (x > 0.5 && y > 0.5)
    {
      primary_value = ValueType{ 1.5, 0, 0, 0, 1.5 };
    }
    else if (x < 0.5 && y > 0.5)
    {
      primary_value = ValueType{ 0.5323, 1.206, 0, 0, 0.3 };
    }
    else if (x < 0.5 && y < 0.5)
    {
      primary_value = ValueType{ 0.138, 1.206, 1.206, 0, 0.029 };
    }
    else // if (x > 0.5 && y < 0.5)
    {
      primary_value = ValueType{ 0.5323, 0, 1.206, 0, 0.3 };
    }

    return PrimaryToConservation(primary_value);
  }

  VTKM_EXEC ValueType PrimaryToConservation(const ValueType& primary_value) const
  {
    FloatDefault gam = 1.4, gamm1 = gam - 1;
    ValueType value;
    auto r = primary_value[0];
    auto u = primary_value[1];
    auto v = primary_value[2];
    auto w = primary_value[3];
    auto p = primary_value[4];

    value[0] = r;
    value[1] = r * u;
    value[2] = r * v;
    value[3] = r * w;
    value[4] = p / gamm1 + 0.5 * r * (u * u + v * v);
    return value;
  }
};

RegisterObject(Riemann2D);
Riemann2D::Riemann2D(const Configuration& cfg)
  : EulerEquation(cfg)
{
}


struct ComputeResidual_Riemann2D : vtkm::worklet::WorkletCellNeighborhood
{
  VTKM_CONT
  ComputeResidual_Riemann2D(FloatDefault t)
    : _t(t)
  {
  }

  using ControlSignature =
    void(CellSetIn, FieldInNeighborhood, FieldOut Rhs, FieldIn centroid, FieldIn spacing);
  using ExecutionSignature = void(Boundary, _2, _3, _4, _5);

  template<typename NeighIn, typename RhsType>
  VTKM_EXEC void operator()(const vtkm::exec::BoundaryState& boundary,
                            const NeighIn& value,
                            RhsType& rhs,
                            const Vec3f& centroid,
                            const Vec3f& spacing) const
  {
    Vec3f normal{ 0, 0, 0 };
    DataPack<RhsType::NUM_COMPONENTS> dl, dr;
    dl.uh = value.Get(0, 0, 0);
    dl.InvisTerm();
    {
      dl.uh = value.Get(0, 0, 0);
      dr.uh = value.Get(-1, 0, 0);
      RhsType lflux, rflux;
      if (boundary.IJK[0] == 0)
      {
        Vec3f p = centroid;
        p[0] = p[0] - 0.5 * spacing[0];
        BoundaryCondition(dl, normal, dr.uh);
      }
      dr.InvisTerm();

      normal = Vec3f{ -1, 0, 0 };
      RiemannFlux(dl, dr, normal, lflux);

      dr.uh = value.Get(1, 0, 0);
      if (boundary.IJK[0] == boundary.PointDimensions[0] - 1)
      {
        Vec3f p = centroid;
        p[0] = p[0] + 0.5 * spacing[0];
        dr.uh = value.Get(0, 0, 0);
        //dr.uh = _exact_value(_t, p);
        BoundaryCondition(dl, normal, dr.uh);
      }
      dr.InvisTerm();
      normal = Vec3f{ 1, 0, 0 };
      RiemannFlux(dl, dr, normal, rflux);
      rhs = (rflux + lflux) / spacing[0];
    }

    {
      dr.uh = value.Get(0, -1, 0);
      RhsType lflux, rflux;
      if (boundary.IJK[1] == 0)
      {
        Vec3f p = centroid;
        p[1] = p[1] - 0.5 * spacing[1];
        BoundaryCondition(dl, normal, dr.uh);
      }

      dr.InvisTerm();

      normal = Vec3f{ 0, -1, 0 };
      RiemannFlux(dl, dr, normal, lflux);

      dr.uh = value.Get(0, 1, 0);
      if (boundary.IJK[1] == boundary.PointDimensions[1] - 1)
      {
        Vec3f p = centroid;
        p[1] = p[1] + 0.5 * spacing[1];
        BoundaryCondition(dl, normal, dr.uh);
      }
      dr.InvisTerm();

      normal = Vec3f{ 0, 1, 0 };
      RiemannFlux(dl, dr, normal, rflux);
      rhs += (rflux + lflux) / spacing[1];
    }
  }

  template<typename DataPackType, typename VType>
  VTKM_EXEC void RiemannFlux(const DataPackType& dl,
                             const DataPackType& dr,
                             const Vec3f& normal,
                             VType& riemann_flux) const
  {
    for (IdComponent eq = 0; eq < VType::NUM_COMPONENTS; eq++)
    {
      riemann_flux[eq] =
        vtkm::Dot(dl.inv_term[eq] + dr.inv_term[eq], normal) / 2 + 0.3 * (dl.uh[eq] - dr.uh[eq]);
    }
  }

  template<typename DataPackType, typename VType>
  VTKM_EXEC void BoundaryCondition(const DataPackType& dl, const Vec3f& normal, VType& ur)
  {
    auto& ul = dl.uh;
    Vec3f mom{ ul[1], ul[2], ul[3] };
    Real vn = vtkm::Dot(mom, normal);
    ur[0] = ul[0];
    ur[1] = ul[1] - 2 * vn * normal[0];
    ur[2] = ul[2] - 2 * vn * normal[1];
    ur[3] = ul[3] - 2 * vn * normal[2];
    ur[4] = ul[4];
  }
  FloatDefault _t;
};

void Riemann2D::InitialCondition() {
  vtkm::cont::Timer timer;
  timer.Start();
  vtkm::cont::Invoker invoker;
  invoker(Riemann2DIC{}, _mesh.GetCellSet(), _solution, _mesh.GetCoordinate());

  console::Success("初始条件设置完成", timer.GetElapsedTime(), "s");
}
